/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.io.kafka;



import com.fasterxml.jackson.annotation.JsonProperty;

import com.fasterxml.jackson.databind.ObjectMapper;

import com.bff.gaia.unified.sdk.coders.BigEndianLongCoder;

import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.coders.KvCoder;

import com.bff.gaia.unified.sdk.metrics.Counter;

import com.bff.gaia.unified.sdk.metrics.Metrics;

import com.bff.gaia.unified.sdk.metrics.SinkMetrics;

import com.bff.gaia.unified.sdk.state.BagState;

import com.bff.gaia.unified.sdk.state.StateSpec;

import com.bff.gaia.unified.sdk.state.StateSpecs;

import com.bff.gaia.unified.sdk.state.ValueState;

import com.bff.gaia.unified.sdk.transforms.DoFn;

import com.bff.gaia.unified.sdk.transforms.GroupByKey;

import com.bff.gaia.unified.sdk.transforms.PTransform;

import com.bff.gaia.unified.sdk.transforms.ParDo;

import com.bff.gaia.unified.sdk.transforms.windowing.AfterPane;

import com.bff.gaia.unified.sdk.transforms.windowing.GlobalWindows;

import com.bff.gaia.unified.sdk.transforms.windowing.Repeatedly;

import com.bff.gaia.unified.sdk.transforms.windowing.Window;

import com.bff.gaia.unified.sdk.values.KV;

import com.bff.gaia.unified.sdk.values.PCollection;

import com.bff.gaia.unified.sdk.values.TimestampedValue;

import com.bff.gaia.unified.vendor.guava.com.google.common.base.MoreObjects;

import com.bff.gaia.unified.vendor.guava.com.google.common.cache.*;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableList;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableMap;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Iterators;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Lists;

import org.apache.kafka.clients.consumer.Consumer;

import org.apache.kafka.clients.consumer.ConsumerConfig;

import org.apache.kafka.clients.consumer.OffsetAndMetadata;

import org.apache.kafka.clients.producer.*;

import org.apache.kafka.common.KafkaException;

import org.apache.kafka.common.TopicPartition;

import org.apache.kafka.common.serialization.ByteArrayDeserializer;



import org.joda.time.DateTimeUtils;

import org.joda.time.DateTimeZone;

import org.joda.time.format.DateTimeFormat;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import java.io.IOException;

import java.util.*;

import java.util.concurrent.*;



import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions.*;



/**

 * Exactly-once sink transform for Kafka. See {@link KafkaIO} for user visible documentation and

 * example usage.

 */

class KafkaExactlyOnceSink<K, V>

    extends PTransform<PCollection<ProducerRecord<K, V>>, PCollection<Void>> {



  // Dataflow ensures at-least once processing for side effects like sinks. In order to provide

  // exactly-once semantics, a sink needs to be idempotent or it should avoid writing records

  // that have already been written. This snk does the latter. All the the records are ordered

  // across a fixed number of shards and records in each shard are written in order. It drops

  // any records that are already written and buffers those arriving out of order.

  //

  // Exactly once sink involves two shuffles of the records:

  //    A : Assign a shard ---> B : Assign sequential ID ---> C : Write to Kafka in order

  //

  // Processing guarantees also require deterministic processing within user transforms.

  // Here, that requires order of the records committed to Kafka by C should not be affected by

  // restarts in C and its upstream stages.

  //

  // A : Assigns a random shard for message. Note that there are no ordering guarantees for

  //     writing user records to Kafka. User can still control partitioning among topic

  //     partitions as with regular sink (of course, there are no ordering guarantees in

  //     regular Kafka sink either).

  // B : Assigns an id sequentially for each messages within a shard.

  // C : Writes each shard to Kafka in sequential id order. In Dataflow, when C sees a record

  //     and id, it implies that record and the associated id are checkpointed to persistent

  //     storage and this record will always have same id, even in retries.

  //     Exactly-once semantics are achieved by writing records in the strict order of

  //     these check-pointed sequence ids.

  //

  // Parallelism for B and C is fixed to 'numShards', which defaults to number of partitions

  // for the topic. A few reasons for that:

  //  - B & C implement their functionality using per-key state. Shard id makes it independent

  //    of cardinality of user key.

  //  - We create one producer per shard, and its 'transactional id' is based on shard id. This

  //    requires that number of shards to be finite. This also helps with batching. and avoids

  //    initializing producers and transactions.

  //  - Most importantly, each of sharded writers stores 'next message id' in partition

  //    metadata, which is committed atomically with Kafka transactions. This is critical

  //    to handle retries of C correctly. Initial testing showed number of shards could be

  //    larger than number of partitions for the topic.

  //

  // Number of shards can change across multiple runs of a pipeline (job upgrade in Dataflow).

  //



  private static final Logger LOG = LoggerFactory.getLogger(KafkaExactlyOnceSink.class);

  private static final String METRIC_NAMESPACE = "KafkaExactlyOnceSink";



  private final KafkaIO.WriteRecords<K, V> spec;



  static void ensureEOSSupport() {

    checkArgument(

        ProducerSpEL.supportsTransactions(),

        "%s %s",

        "This version of Kafka client does not support transactions required to support",

        "exactly-once semantics. Please use Kafka client version 0.11 or newer.");

  }



  KafkaExactlyOnceSink(KafkaIO.WriteRecords<K, V> spec) {

    this.spec = spec;

  }



  @Override

  public PCollection<Void> expand(PCollection<ProducerRecord<K, V>> input) {



    int numShards = spec.getNumShards();

    if (numShards <= 0) {

      try (Consumer<?, ?> consumer = openConsumer(spec)) {

        numShards = consumer.partitionsFor(spec.getTopic()).size();

        LOG.info(

            "Using {} shards for exactly-once writer, matching number of partitions "

                + "for topic '{}'",

            numShards,

            spec.getTopic());

      }

    }

    checkState(numShards > 0, "Could not set number of shards");



    return input

        .apply(

            Window.<ProducerRecord<K, V>>into(new GlobalWindows()) // Everything into global window.

                .triggering(Repeatedly.forever(AfterPane.elementCountAtLeast(1)))

                .discardingFiredPanes())

        .apply(

            String.format("Shuffle across %d shards", numShards),

            ParDo.of(new Reshard<>(numShards)))

        .apply("Persist sharding", GroupByKey.create())

        .apply("Assign sequential ids", ParDo.of(new Sequencer<>()))

        .apply("Persist ids", GroupByKey.create())

        .apply(

            String.format("Write to Kafka topic '%s'", spec.getTopic()),

            ParDo.of(new ExactlyOnceWriter<>(spec, input.getCoder())));

  }



  /** Shuffle messages assigning each randomly to a shard. */

  private static class Reshard<K, V>

      extends DoFn<ProducerRecord<K, V>, KV<Integer, TimestampedValue<ProducerRecord<K, V>>>> {



    private final int numShards;

    private transient int shardId;



    Reshard(int numShards) {

      this.numShards = numShards;

    }



    @Setup

    public void setup() {

      shardId = ThreadLocalRandom.current().nextInt(numShards);

    }



    @ProcessElement

    public void processElement(ProcessContext ctx) {

      shardId = (shardId + 1) % numShards; // round-robin among shards.

      ctx.output(KV.of(shardId, TimestampedValue.of(ctx.element(), ctx.timestamp())));

    }

  }



  private static class Sequencer<K, V>

      extends DoFn<

          KV<Integer, Iterable<TimestampedValue<ProducerRecord<K, V>>>>,

          KV<Integer, KV<Long, TimestampedValue<ProducerRecord<K, V>>>>> {



    private static final String NEXT_ID = "nextId";



    @StateId(NEXT_ID)

    private final StateSpec<ValueState<Long>> nextIdSpec = StateSpecs.value();



    @ProcessElement

    public void processElement(@StateId(NEXT_ID) ValueState<Long> nextIdState, ProcessContext ctx) {

      long nextId = MoreObjects.firstNonNull(nextIdState.read(), 0L);

      int shard = ctx.element().getKey();

      for (TimestampedValue<ProducerRecord<K, V>> value : ctx.element().getValue()) {

        ctx.output(KV.of(shard, KV.of(nextId, value)));

        nextId++;

      }

      nextIdState.write(nextId);

    }

  }



  private static class ExactlyOnceWriter<K, V>

      extends DoFn<KV<Integer, Iterable<KV<Long, TimestampedValue<ProducerRecord<K, V>>>>>, Void> {



    private static final String NEXT_ID = "nextId";

    private static final String MIN_BUFFERED_ID = "minBufferedId";

    private static final String OUT_OF_ORDER_BUFFER = "outOfOrderBuffer";

    private static final String WRITER_ID = "writerId";



    // Not sure of a good limit. This applies only for large bundles.

    private static final int MAX_RECORDS_PER_TXN = 1000;

    private static final ObjectMapper JSON_MAPPER = new ObjectMapper();



    @StateId(NEXT_ID)

    private final StateSpec<ValueState<Long>> sequenceIdSpec = StateSpecs.value();



    @StateId(MIN_BUFFERED_ID)

    private final StateSpec<ValueState<Long>> minBufferedIdSpec = StateSpecs.value();



    @StateId(OUT_OF_ORDER_BUFFER)

    private final StateSpec<BagState<KV<Long, TimestampedValue<ProducerRecord<K, V>>>>>

        outOfOrderBufferSpec;

    // A random id assigned to each shard. Helps with detecting when multiple jobs are mistakenly

    // started with same groupId used for storing state on Kafka side, including the case where

    // a job is restarted with same groupId, but the metadata from previous run was not cleared.

    // Better to be safe and error out with a clear message.

    @StateId(WRITER_ID)

    private final StateSpec<ValueState<String>> writerIdSpec = StateSpecs.value();



    private final KafkaIO.WriteRecords<K, V> spec;



    // Metrics

    private final Counter elementsWritten = SinkMetrics.elementsWritten();

    // Elements buffered due to out of order arrivals.

    private final Counter elementsBuffered = Metrics.counter(METRIC_NAMESPACE, "elementsBuffered");

    private final Counter numTransactions = Metrics.counter(METRIC_NAMESPACE, "numTransactions");



    ExactlyOnceWriter(KafkaIO.WriteRecords<K, V> spec, Coder<ProducerRecord<K, V>> elemCoder) {

      this.spec = spec;

      this.outOfOrderBufferSpec =

          StateSpecs.bag(KvCoder.of(BigEndianLongCoder.of(), TimestampedValue.TimestampedValueCoder.of(elemCoder)));

    }



    @Setup

    public void setup() {

      // This is on the worker. Ensure the runtime version is till compatible.

      KafkaExactlyOnceSink.ensureEOSSupport();

    }



    // Futures ignored as exceptions will be flushed out in the commitTxn

    @SuppressWarnings("FutureReturnValueIgnored")

    @RequiresStableInput

    @ProcessElement

    public void processElement(

        @StateId(NEXT_ID) ValueState<Long> nextIdState,

        @StateId(MIN_BUFFERED_ID) ValueState<Long> minBufferedIdState,

        @StateId(OUT_OF_ORDER_BUFFER)

            BagState<KV<Long, TimestampedValue<ProducerRecord<K, V>>>> oooBufferState,

        @StateId(WRITER_ID) ValueState<String> writerIdState,

        ProcessContext ctx)

        throws IOException {



      int shard = ctx.element().getKey();



      minBufferedIdState.readLater();

      long nextId = MoreObjects.firstNonNull(nextIdState.read(), 0L);

      long minBufferedId = MoreObjects.firstNonNull(minBufferedIdState.read(), Long.MAX_VALUE);



      ShardWriterCache<K, V> cache =

          (ShardWriterCache<K, V>) CACHE_BY_GROUP_ID.getUnchecked(spec.getSinkGroupId());

      ShardWriter<K, V> writer = cache.removeIfPresent(shard);

      if (writer == null) {

        writer = initShardWriter(shard, writerIdState, nextId);

      }



      long committedId = writer.committedId;



      if (committedId >= nextId) {

        // This is a retry of an already committed batch.

        LOG.info(

            "{}: committed id {} is ahead of expected {}. {} records will be dropped "

                + "(these are already written).",

            shard,

            committedId,

            nextId - 1,

            committedId - nextId + 1);

        nextId = committedId + 1;

      }



      try {

        writer.beginTxn();

        int txnSize = 0;



        // Iterate in recordId order. The input iterator could be mostly sorted.

        // There might be out of order messages buffered in earlier iterations. These

        // will get merged if and when minBufferedId matches nextId.



        Iterator<KV<Long, TimestampedValue<ProducerRecord<K, V>>>> iter =

            ctx.element().getValue().iterator();



        while (iter.hasNext()) {

          KV<Long, TimestampedValue<ProducerRecord<K, V>>> kv = iter.next();

          long recordId = kv.getKey();



          if (recordId < nextId) {

            LOG.info(

                "{}: dropping older record {}. Already committed till {}",

                shard,

                recordId,

                committedId);

            continue;

          }



          if (recordId > nextId) {

            // Out of order delivery. Should be pretty rare (what about in a batch pipeline?)

            LOG.info(

                "{}: Saving out of order record {}, next record id to be written is {}",

                shard,

                recordId,

                nextId);



            // checkState(recordId - nextId < 10000, "records are way out of order");



            oooBufferState.add(kv);

            minBufferedId = Math.min(minBufferedId, recordId);

            minBufferedIdState.write(minBufferedId);

            elementsBuffered.inc();

            continue;

          }



          // recordId and nextId match. Finally write the record.



          writer.sendRecord(kv.getValue(), elementsWritten);

          nextId++;



          if (++txnSize >= MAX_RECORDS_PER_TXN) {

            writer.commitTxn(recordId, numTransactions);

            txnSize = 0;

            writer.beginTxn();

          }



          if (minBufferedId == nextId) {

            // One or more of the buffered records can be committed now.

            // Read all of them in to memory and sort them. Reading into memory

            // might be problematic in extreme cases. Might need to improve it in future.



            List<KV<Long, TimestampedValue<ProducerRecord<K, V>>>> buffered =

                Lists.newArrayList(oooBufferState.read());

            buffered.sort(new KV.OrderByKey<>());



            LOG.info(

                "{} : merging {} buffered records (min buffered id is {}).",

                shard,

                buffered.size(),

                minBufferedId);



            oooBufferState.clear();

            minBufferedIdState.clear();

            minBufferedId = Long.MAX_VALUE;



            iter =

                Iterators.mergeSorted(

                    ImmutableList.of(iter, buffered.iterator()), new KV.OrderByKey<>());

          }

        }



        writer.commitTxn(nextId - 1, numTransactions);

        nextIdState.write(nextId);



      } catch (ProducerSpEL.UnrecoverableProducerException e) {

        // Producer JavaDoc says these are not recoverable errors and producer should be closed.



        // Close the producer and a new producer will be initialized in retry.

        // It is possible that a rough worker keeps retrying and ends up fencing off

        // active producers. How likely this might be or how well such a scenario is handled

        // depends on the runner. For now we will leave it to upper layers, will need to revisit.



        LOG.warn(

            "{} : closing producer {} after unrecoverable error. The work might have migrated."

                + " Committed id {}, current id {}.",

            writer.shard,

            writer.producerName,

            writer.committedId,

            nextId - 1,

            e);



        writer.producer.close();

        writer = null; // No need to cache it.

        throw e;

      } finally {

        if (writer != null) {

          cache.insert(shard, writer);

        }

      }

    }



    private static class ShardMetadata {



      @JsonProperty("seq")

      public final long sequenceId;



      @JsonProperty("id")

      public final String writerId;



      private ShardMetadata() { // for json deserializer

        sequenceId = -1;

        writerId = null;

      }



      ShardMetadata(long sequenceId, String writerId) {

        this.sequenceId = sequenceId;

        this.writerId = writerId;

      }

    }



    /** A wrapper around Kafka producer. One for each of the shards. */

    private static class ShardWriter<K, V> {



      private final int shard;

      private final String writerId;

      private final Producer<K, V> producer;

      private final String producerName;

      private final KafkaIO.WriteRecords<K, V> spec;

      private long committedId;



      ShardWriter(

          int shard,

          String writerId,

          Producer<K, V> producer,

          String producerName,

          KafkaIO.WriteRecords<K, V> spec,

          long committedId) {

        this.shard = shard;

        this.writerId = writerId;

        this.producer = producer;

        this.producerName = producerName;

        this.spec = spec;

        this.committedId = committedId;

      }



      void beginTxn() {

        ProducerSpEL.beginTransaction(producer);

      }



      Future<RecordMetadata> sendRecord(

          TimestampedValue<ProducerRecord<K, V>> record, Counter sendCounter) {

        try {

          Long timestampMillis =

              spec.getPublishTimestampFunction() != null

                  ? spec.getPublishTimestampFunction()

                      .getTimestamp(record.getValue(), record.getTimestamp())

                      .getMillis()

                  : null;



          Future<RecordMetadata> result =

              producer.send(

                  new ProducerRecord<>(

                      spec.getTopic(),

                      null,

                      timestampMillis,

                      record.getValue().key(),

                      record.getValue().value()));

          sendCounter.inc();

          return result;

        } catch (KafkaException e) {

          ProducerSpEL.abortTransaction(producer);

          throw e;

        }

      }



      void commitTxn(long lastRecordId, Counter numTransactions) throws IOException {

        try {

          // Store id in consumer group metadata for the partition.

          // NOTE: Kafka keeps this metadata for 24 hours since the last update. This limits

          // how long the pipeline could be down before resuming it. It does not look like

          // this TTL can be adjusted (asked about it on Kafka users list).

          ProducerSpEL.sendOffsetsToTransaction(

              producer,

              ImmutableMap.of(

                  new TopicPartition(spec.getTopic(), shard),

                  new OffsetAndMetadata(

                      0L,

                      JSON_MAPPER.writeValueAsString(new ShardMetadata(lastRecordId, writerId)))),

              spec.getSinkGroupId());

          ProducerSpEL.commitTransaction(producer);



          numTransactions.inc();

          LOG.debug("{} : committed {} records", shard, lastRecordId - committedId);



          committedId = lastRecordId;

        } catch (KafkaException e) {

          ProducerSpEL.abortTransaction(producer);

          throw e;

        }

      }

    }



    private ShardWriter<K, V> initShardWriter(

        int shard, ValueState<String> writerIdState, long nextId) throws IOException {



      String producerName = String.format("producer_%d_for_%s", shard, spec.getSinkGroupId());

      Producer<K, V> producer = initializeExactlyOnceProducer(spec, producerName);



      // Fetch latest committed metadata for the partition (if any). Checks committed sequence ids.

      try {



        String writerId = writerIdState.read();



        OffsetAndMetadata committed;



        try (Consumer<?, ?> consumer = openConsumer(spec)) {

          committed = consumer.committed(new TopicPartition(spec.getTopic(), shard));

        }



        long committedSeqId = -1;



        if (committed == null || committed.metadata() == null || committed.metadata().isEmpty()) {

          checkState(

              nextId == 0 && writerId == null,

              "State exists for shard %s (nextId %s, writerId '%s'), but there is no state "

                  + "stored with Kafka topic '%s' group id '%s'",

              shard,

              nextId,

              writerId,

              spec.getTopic(),

              spec.getSinkGroupId());



          writerId =

              String.format(

                  "%X - %s",

                  new Random().nextInt(Integer.MAX_VALUE),

                  DateTimeFormat.forPattern("yyyy-MM-dd HH:mm:ss")

                      .withZone(DateTimeZone.UTC)

                      .print(DateTimeUtils.currentTimeMillis()));

          writerIdState.write(writerId);

          LOG.info("Assigned writer id '{}' to shard {}", writerId, shard);



        } else {

          ShardMetadata metadata = JSON_MAPPER.readValue(committed.metadata(), ShardMetadata.class);



          checkNotNull(metadata.writerId);



          if (writerId == null) {

            // a) This might be a restart of the job from scratch, in which case metatdata

            // should be ignored and overwritten with new one.

            // b) This job might be started with an incorrect group id which is an error.

            // c) There is an extremely small chance that this is a retry of the first bundle

            // where metatdate was committed to Kafka but the bundle results were not committed

            // in Unified, in which case it should be treated as correct metadata.

            // How can we tell these three cases apart? Be safe and throw an exception.

            //

            // We could let users explicitly an option to override the existing metadata.

            //

            throw new IllegalStateException(

                String.format(

                    "Kafka metadata exists for shard %s, but there is no stored state for it. "

                        + "This mostly indicates groupId '%s' is used else where or in earlier runs. "

                        + "Try another group id. Metadata for this shard on Kafka : '%s'",

                    shard, spec.getSinkGroupId(), committed.metadata()));

          }



          checkState(

              writerId.equals(metadata.writerId),

              "Writer ids don't match. This is mostly a unintended misuse of groupId('%s')."

                  + "Unified '%s', Kafka '%s'",

              spec.getSinkGroupId(),

              writerId,

              metadata.writerId);



          committedSeqId = metadata.sequenceId;



          checkState(

              committedSeqId >= (nextId - 1),

              "Committed sequence id can not be lower than %s, partition metadata : %s",

              nextId - 1,

              committed.metadata());

        }



        LOG.info(

            "{} : initialized producer {} with committed sequence id {}",

            shard,

            producerName,

            committedSeqId);



        return new ShardWriter<>(shard, writerId, producer, producerName, spec, committedSeqId);



      } catch (IOException e) {

        producer.close();

        throw e;

      }

    }



    /**

     * A wrapper around guava cache to provide insert()/remove() semantics. A ShardWriter will be

     * closed if it is stays in cache for more than 1 minute, i.e. not used inside

     * KafkaExactlyOnceSink DoFn for a minute.

     */

    private static class ShardWriterCache<K, V> {



      static final ScheduledExecutorService SCHEDULED_CLEAN_UP_THREAD =

          Executors.newSingleThreadScheduledExecutor();



      static final int CLEAN_UP_CHECK_INTERVAL_MS = 10 * 1000;

      static final int IDLE_TIMEOUT_MS = 60 * 1000;



      private final Cache<Integer, ShardWriter<K, V>> cache;



      // Exceptions arising from the cache cleanup are ignored

      @SuppressWarnings("FutureReturnValueIgnored")

      ShardWriterCache() {

        this.cache =

            CacheBuilder.newBuilder()

                .expireAfterWrite(IDLE_TIMEOUT_MS, TimeUnit.MILLISECONDS)

                .<Integer, ShardWriter<K, V>>removalListener(

                    notification -> {

                      if (notification.getCause() != RemovalCause.EXPLICIT) {

                        ShardWriter writer = notification.getValue();

                        LOG.info(

                            "{} : Closing idle shard writer {} after 1 minute of idle time.",

                            writer.shard,

                            writer.producerName);

                        writer.producer.close();

                      }

                    })

                .build();



        // run cache.cleanUp() every 10 seconds.

        SCHEDULED_CLEAN_UP_THREAD.scheduleAtFixedRate(

            cache::cleanUp,

            CLEAN_UP_CHECK_INTERVAL_MS,

            CLEAN_UP_CHECK_INTERVAL_MS,

            TimeUnit.MILLISECONDS);

      }



      ShardWriter<K, V> removeIfPresent(int shard) {

        return cache.asMap().remove(shard);

      }



      void insert(int shard, ShardWriter<K, V> writer) {

        ShardWriter<K, V> existing = cache.asMap().putIfAbsent(shard, writer);

        checkState(

            existing == null, "Unexpected multiple instances of writers for shard %s", shard);

      }

    }



    // One cache for each sink (usually there is only one sink per pipeline)

    private static final LoadingCache<String, ShardWriterCache<?, ?>> CACHE_BY_GROUP_ID =

        CacheBuilder.newBuilder()

            .build(

                new CacheLoader<String, ShardWriterCache<?, ?>>() {

                  @Override

                  public ShardWriterCache<?, ?> load(String key) throws Exception {

                    return new ShardWriterCache<>();

                  }

                });

  }



  /**

   * Opens a generic consumer that is mainly meant for metadata operations like fetching number of

   * partitions for a topic rather than for fetching messages.

   */

  private static Consumer<?, ?> openConsumer(KafkaIO.WriteRecords<?, ?> spec) {

    return spec.getConsumerFactoryFn()

        .apply(

            ImmutableMap.of(

                ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG,

                spec.getProducerConfig().get(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG),

                ConsumerConfig.GROUP_ID_CONFIG,

                spec.getSinkGroupId(),

                ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG,

                ByteArrayDeserializer.class,

                ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG,

                ByteArrayDeserializer.class));

  }



  private static <K, V> Producer<K, V> initializeExactlyOnceProducer(

	  KafkaIO.WriteRecords<K, V> spec, String producerName) {



    Map<String, Object> producerConfig = new HashMap<>(spec.getProducerConfig());

    producerConfig.putAll(

        ImmutableMap.of(

            ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG,

            spec.getKeySerializer(),

            ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG,

            spec.getValueSerializer(),

            ProducerSpEL.ENABLE_IDEMPOTENCE_CONFIG,

            true,

            ProducerSpEL.TRANSACTIONAL_ID_CONFIG,

            producerName));



    Producer<K, V> producer =

        spec.getProducerFactoryFn() != null

            ? spec.getProducerFactoryFn().apply(producerConfig)

            : new KafkaProducer<>(producerConfig);



    ProducerSpEL.initTransactions(producer);

    return producer;

  }

}